# -*- coding: utf-8 -*-
"""
JSON数据训练脚本（集成项目架构）

功能亮点：
1. 复用项目数据库配置
2. 自动分块处理大数据
3. 实时进度监控
4. 智能错误恢复
"""

import json
import time
import sys
import os
import signal
from pathlib import Path
from tqdm import tqdm
import requests
from typing import List, Dict
from loguru import logger

# 添加项目根目录到 Python 路径
current_dir = Path(__file__).resolve().parent
project_root = current_dir.parent.parent
sys.path.insert(0, str(project_root))

# 全局变量用于优雅退出
interrupted = False

def signal_handler(sig, frame):
    """处理中断信号"""
    global interrupted
    interrupted = True
    logger.warning("\n⚠️ 接收到中断信号，正在安全退出...")
    logger.warning("请等待当前批次处理完成...")

class JSONTrainer:
    def __init__(self):
        
        # 接口配置
        self.api_endpoint = "http://127.0.0.1:6011/bailian/v1/chat/chat_with_own_pet"
        
        # 从环境变量读取 Token（安全性更好）
        default_token = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIyMDI1MDYyMzAwMSIsImV4cCI6MTc2Mjk0NDkxMX0.PlxnX1lIBCCzL-34tl5tbwkfZGG2WXIIuzRyXdz1wIA"
        self.headers = {
            "Content-Type": "application/json",
            "Authorization": os.getenv("TRAINING_TOKEN", default_token)
        }
        
        # 训练配置
        self.chunk_size = 10  # 每批处理量
        self.retry_limit = 3  # 失败重试次数
        self.request_interval = 1  # 请求间隔(秒)
        
        # 失败记录
        self.failed_records = []

    def load_json_data(self, file_path: str) -> List[Dict]:
        """ 安全加载JSON数据 """
        try:
            logger.info(f"正在加载数据文件: {file_path}")
            with open(file_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
            
            if not isinstance(data, list):
                raise ValueError(f"数据格式错误：期望列表，实际为 {type(data).__name__}")
            
            logger.success(f"✅ 成功加载 {len(data)} 条记录")
            return data
            
        except FileNotFoundError:
            logger.error(f"❌ 文件不存在: {file_path}")
            raise
        except json.JSONDecodeError as e:
            logger.error(f"❌ JSON格式错误: {str(e)}")
            raise
        except Exception as e:
            logger.error(f"❌ 加载失败: {type(e).__name__} - {str(e)}")
            raise

    def validate_record(self, record: Dict) -> bool:
        """ 数据验证 """
        required_fields = ['user_id', 'pet_query', 'query']
        
        # 检查必需字段是否存在
        if not all(field in record for field in required_fields):
            missing_fields = [f for f in required_fields if f not in record]
            logger.warning(f"记录缺少必需字段: {missing_fields} | 记录: {record}")
            return False
        
        # 检查 user_id 和 query 是否为空
        if not record.get('user_id') or not str(record['user_id']).strip():
            logger.warning(f"user_id 为空 | 记录: {record}")
            return False
            
        if not record.get('query') or not str(record['query']).strip():
            logger.warning(f"query 为空 | 记录: {record}")
            return False
        
        return True

    def process_chunk(self, chunk: List[Dict]) -> Dict:
        """ 处理数据块 """
        stats = {'success': 0, 'failed': 0, 'skipped': 0}
        
        for record in chunk:
            if not self.validate_record(record):
                stats['skipped'] += 1
                continue
                
            payload = {
                "uid": record['user_id'],
                "pet_query": record['pet_query'] or "",
                "query": record['query'],
                "nickname": "训练员",
                "character": "严谨",
                "whether_train": True,
                "stream": False
            }
            
            success = False
            last_error = None
            
            for attempt in range(self.retry_limit):
                try:
                    response = requests.post(
                        self.api_endpoint,
                        json=payload,
                        headers=self.headers,
                        timeout=30
                    )
                    
                    if response.status_code == 200:
                        stats['success'] += 1
                        success = True
                        logger.debug(f"✅ 训练成功 | UID: {record['user_id']}")
                        break
                    else:
                        last_error = f"HTTP {response.status_code}: {response.text[:100]}"
                        logger.warning(f"尝试 {attempt+1}/{self.retry_limit} | UID: {record['user_id']} | {last_error}")
                        
                except requests.Timeout:
                    last_error = "请求超时"
                    logger.warning(f"尝试 {attempt+1}/{self.retry_limit} | UID: {record['user_id']} | {last_error}")
                    
                except requests.ConnectionError:
                    last_error = "连接失败，请检查服务是否启动"
                    logger.warning(f"尝试 {attempt+1}/{self.retry_limit} | UID: {record['user_id']} | {last_error}")
                    
                except Exception as e:
                    last_error = f"{type(e).__name__}: {str(e)}"
                    logger.warning(f"尝试 {attempt+1}/{self.retry_limit} | UID: {record['user_id']} | {last_error}")
                
                # 指数退避
                if attempt < self.retry_limit - 1:
                    wait_time = 2 ** attempt
                    time.sleep(wait_time)
            
            if not success:
                stats['failed'] += 1
                record['_error'] = last_error  # 记录失败原因
                self.failed_records.append(record)
                logger.error(f"❌ 记录永久失败 | UID: {record['user_id']} | 原因: {last_error}")
            
            time.sleep(self.request_interval)
        
        return stats

    def run_training(self, file_path: str):
        """ 执行训练流程 """
        start_time = time.time()
        logger.info("=" * 60)
        logger.info("🚀 启动JSON数据训练")
        logger.info("=" * 60)
        
        # 加载数据
        try:
            all_data = self.load_json_data(file_path)
            total = len(all_data)
        except Exception as e:
            logger.error(f"训练终止：数据加载失败")
            return
        
        # 显示训练配置
        logger.info(f"训练配置:")
        logger.info(f"  接口地址: {self.api_endpoint}")
        logger.info(f"  总记录数: {total}")
        logger.info(f"  批次大小: {self.chunk_size}")
        logger.info(f"  重试次数: {self.retry_limit}")
        logger.info(f"  请求间隔: {self.request_interval}秒")
        logger.info("=" * 60)

        # 分块处理
        progress = tqdm(total=total, desc="训练进度", unit="条")
        global_stats = {'success': 0, 'failed': 0, 'skipped': 0}
        
        for i in range(0, total, self.chunk_size):
            # 检查是否需要中断
            global interrupted
            if interrupted:
                logger.warning("⚠️ 训练已被用户中断")
                progress.close()
                break
            
            chunk = all_data[i:i+self.chunk_size]
            chunk_stats = self.process_chunk(chunk)
            
            global_stats['success'] += chunk_stats['success']
            global_stats['failed'] += chunk_stats['failed']
            global_stats['skipped'] += chunk_stats['skipped']
            
            progress.update(len(chunk))
            
            # 计算成功率（避免除零错误）
            total_processed = global_stats['success'] + global_stats['failed']
            success_rate = global_stats['success'] / total_processed if total_processed > 0 else 0
            
            progress.set_postfix({
                "成功": global_stats['success'],
                "失败": global_stats['failed'],
                "跳过": global_stats['skipped'],
                "成功率": f"{success_rate:.2%}"
            })
        
        if not interrupted:
            progress.close()
        
        # 计算训练时长
        duration = time.time() - start_time
        
        # 生成报告
        self.generate_report(global_stats, duration)
        
        # 保存失败记录
        if self.failed_records:
            self.save_failed_records()

    def save_failed_records(self):
        """保存失败的记录到文件，便于后续重试"""
        failed_file = project_root / "datas_failed.json"
        try:
            with open(failed_file, 'w', encoding='utf-8') as f:
                json.dump(self.failed_records, f, ensure_ascii=False, indent=4)
            logger.warning(f"⚠️ 失败记录已保存到: {failed_file}")
            logger.info(f"可以修复问题后，将该文件重命名为 datas.json 重新训练")
        except Exception as e:
            logger.error(f"保存失败记录时出错: {e}")

    def generate_report(self, stats: Dict, duration: float):
        """ 生成训练报告 """
        total_processed = stats['success'] + stats['failed']
        total_records = total_processed + stats['skipped']
        success_rate = stats['success'] / total_processed if total_processed > 0 else 0
        
        # 格式化时长
        hours = int(duration // 3600)
        minutes = int((duration % 3600) // 60)
        seconds = int(duration % 60)
        
        if hours > 0:
            duration_str = f"{hours}小时{minutes}分{seconds}秒"
        elif minutes > 0:
            duration_str = f"{minutes}分{seconds}秒"
        else:
            duration_str = f"{seconds}秒"
        
        logger.info("")
        logger.info("=" * 60)
        logger.info("🎉 训练完成!")
        logger.info("=" * 60)
        logger.info(f"📊 统计信息:")
        logger.info(f"  总记录数:    {total_records}")
        logger.info(f"  ✅ 成功:      {stats['success']} ({success_rate:.2%})")
        logger.info(f"  ❌ 失败:      {stats['failed']}")
        logger.info(f"  ⚠️  跳过:      {stats['skipped']}")
        logger.info(f"  ⏱️  总耗时:    {duration_str}")
        logger.info(f"  ⚡ 平均速度:  {total_processed/duration:.2f} 条/秒" if duration > 0 else "  ⚡ 平均速度:  N/A")
        logger.info("=" * 60)

if __name__ == "__main__":
    # 注册信号处理器（Ctrl+C 优雅退出）
    signal.signal(signal.SIGINT, signal_handler)
    
    # 自动查找 datas.json 文件
    data_file = project_root / "datas.json"
    
    if not data_file.exists():
        logger.error(f"❌ 数据文件不存在: {data_file}")
        logger.info("💡 请先运行 export_training_data.py 导出训练数据")
        logger.info(f"   预期文件路径: {data_file}")
        sys.exit(1)
    
    logger.info(f"📁 数据文件: {data_file}")
    logger.info("💡 提示: 按 Ctrl+C 可以随时安全退出训练")
    logger.info("")
    
    try:
        trainer = JSONTrainer()
        trainer.run_training(str(data_file))
    except KeyboardInterrupt:
        logger.warning("\n⚠️ 训练已被用户中断")
    except Exception as e:
        logger.error(f"\n❌ 训练过程中发生错误: {type(e).__name__} - {str(e)}", exc_info=True)
        sys.exit(1)
    
    if not interrupted:
        logger.success("\n✨ 训练任务全部完成！")